#!/bin/bash
set -uo pipefail
#==============================================================#
# File      :   pg-basebackup
# Desc      :   PostgreSQL base backup script
# Ctime     :   2018-12-06
# Mtime     :   2022-12-17
# Path      :   /pg/bin/pg-basebackup
# Deps      :   pg_basebackup, psql, lz4, openssl, .pgpass
# License   :   AGPLv3 @ https://doc.pgsty.com/about/license
# Copyright :   2018-2025  Ruohang Feng / Vonng (rh@vonng.com)
#==============================================================#
PROG_NAME="$(basename $0)"
PROG_DIR="$(cd $(dirname $0) && pwd)"


#--------------------------------------------------------------#
#                             Usage                            #
#--------------------------------------------------------------#
function usage() {
	cat <<-'EOF'
		NAME
			pg-basebackup  -- make base backup from PostgreSQL instance
		
		SYNOPSIS
			pg-basebackup -sdfeukr
			pg-basebackup --src postgres:/// --dst . --file backup.tar.lz4
		
		DESCRIPTION
			-s, --src, --url
				Backup source URL, optional, "postgres:///" by default
				Note: if password is required, it should be given in url, ENV or .pgpass
		
			-d, --dst, --dir
				Where to put backup files, "/pg/backup" by default
		
			-f, --file
				Overwrite default backup filename, "backup_${tag}_${date}.tar.lz4"
		
			-r, --remove
				.lz4 Files mtime before n minutes ago will be removed, default is 1200 (20hour)
		
			-t, --tag
				Backup file tag, if not set, target cluster_name or local ip address will be used.
				Also used as part of DEFAULT filename
		
			-k, --key
				Encryption key when --encrypt is specified, default key is ${tag}
		
			-u, --upload
				Upload backup files to cloud storage, (need your own implementation)
		
			-e, --encryption
				Encrypt with RC4 using OpenSSL, if not key is specified, tag is used as key
		
			-h, --help
				Print this message
		
		EXAMPLES
			routine backup crontab:
				00 01 * * * /pg/bin/pg-basebackup --encrypt --upload --tag=test --key=<secret> 2>> /pg/log/backup.log
		
			one-time manual backup:
				pg-basebackup -s postgres:/// -d . -f one_time_backup.tar.lz4 -e
		
			extract backup files:
			  backup_dir="/pg/backup"
			  backup_latest=$(ls -t ${backup_dir} | head -n1)
				unlz4 -d -c ${backup_latest} | tar -xC ${DATA_DIR}
				openssl enc -rc4 -d -k ${PASSWORD} -in ${backup_latest} | unlz4 -d -c | tar -xC ${DATA_DIR}
	EOF
}


#--------------------------------------------------------------#
#                             Param                            #
#--------------------------------------------------------------#
export PATH=/usr/pgsql/bin:${PATH}


#--------------------------------------------------------------#
#                             Utils                            #
#--------------------------------------------------------------#
# logger functions
function log_debug() {
	[ -t 2 ] && printf "\033[0;34m[$(date "+%Y-%m-%d %H:%M:%S")][DEBUG] $*\033[0m\n" >&2 ||
		printf "[$(date "+%Y-%m-%d %H:%M:%S")][DEBUG] $*\n" >&2
}
function log_info() {
	[ -t 2 ] && printf "\033[0;32m[$(date "+%Y-%m-%d %H:%M:%S")][INFO] $*\033[0m\n" >&2 ||
		printf "[$(date "+%Y-%m-%d %H:%M:%S")][INFO] $*\n" >&2
}
function log_warn() {
	[ -t 2 ] && printf "\033[0;33m[$(date "+%Y-%m-%d %H:%M:%S")][WARN] $*\033[0m\n" >&2 ||
		printf "[$(date "+%Y-%m-%d %H:%M:%S")][INFO] $*\n" >&2
}
function log_error() {
	[ -t 2 ] && printf "\033[0;31m[$(date "+%Y-%m-%d %H:%M:%S")][ERROR] $*\033[0m\n" >&2 ||
		printf "[$(date "+%Y-%m-%d %H:%M:%S")][INFO] $*\n" >&2
}

# returns 't' on replica, psql access required
function is_in_recovery() {
	local pg_url=$1
	echo $(psql ${pg_url} -AXtwqc "SELECT pg_is_in_recovery();")
}

# get primary IP address (which could be 10.x.x.x or 192.x.x.x)
function local_ip() {
	echo $(/sbin/ifconfig | grep -Eo 'inet (addr:)?([0-9]*\.){3}[0-9]*' |
		grep -v 127.0.0.1 | grep -Eo '(10|192)\.([0-9]*\.){2}[0-9]*' | head -n1)
}

# return cluster_name, psql access required
function get_cluster_name() {
	local pg_url=$1
	local cluster_name=$(psql ${pg_url} -AXtwqc "SHOW cluster_name")
	if [[ -z "${cluster_name}" ]]; then
		echo $(local_ip)
	else
		echo ${cluster_name}
	fi
}

#==============================================================#
#                            Backup                            #
#==============================================================#

#--------------------------------------------------------------#
# Name: make_backup
# Desc: make pg base backup to given path
# Arg1: Postgres URI
# Arg2: Backup filepath
# Arg3: Encrytion key, optional
# Note: if key is provided, encrypt backup with openssl rc4
#--------------------------------------------------------------#
function make_backup() {
	local pg_url=$1
	local backup_path=$2
	local key=${3-''}

	if [[ ! -z "${key}" ]]; then
		# if key is provided, encrypt with rc4 using openssl
		pg_basebackup -w -d ${pg_url} -Xf -Ft -c fast -v -D - | lz4 -q -z | openssl enc -rc4 -k ${key} >"${backup_path}"
		# extract:  openssl enc -rc4 -d -k ${KEY} -in ${BKUP_FILE} | unlz4 -c | tar -xC ${DATA_DIR}
		return $?
	else
		pg_basebackup -w -d ${pg_url} -Xf -Ft -c fast -v -D - | lz4 -q -z >"${backup_path}"
		# extract:  unlz4 ${BKUP_FILE} -c | tar -xC ${DATA_DIR}
		return $?
	fi
}

#--------------------------------------------------------------#
# Name: kill_base_backup
# Desc: kill existing running backup process
#--------------------------------------------------------------#
function kill_base_backup() {
	local pids=$(ps aux | grep pg_basebackup | grep -e "-Xf")
	log_warn "killing basebackup processes ${pids}"

	for pid in ${pids}; do
		log_warn "kill basebackup process: $pid"
		echo $pid | awk '{print $2}' | xargs -n1 kill
		log_info "kill basebackup process ${pid} done"
	done

	log_warn "basebackup processes killed"
}

#--------------------------------------------------------------#
# Name: remove_backup
# Desc: remove old backup files (*.lz4) in given backup dir
# Arg1: backup directory
# Arg2: remove threshhold (minutes, default 1200, i.e 20hour)
#--------------------------------------------------------------#
function remove_backup() {
	# delete *.lz4 file mtime before 20h ago by default
	local backup_dir=$1
	local remove_condition=${2-'1200'}
	remove_condition="-mmin +${remove_condition}"

	log_info "[BKUP] find obsolete backups: find "${backup_dir}/" -maxdepth 1 -type f ${remove_condition} -name 'backup*.lz4'"
	local obsolete_backups="$(find "${backup_dir}/" -maxdepth 1 -type f ${remove_condition} -name 'backup*.lz4')"
	log_warn "[BKUP] remove obsolete backups: ${obsolete_backups}"
	find "${backup_dir}/" -maxdepth 1 -type f -name 'backup_*.lz4' ${remove_condition} -delete
	return $?
}

#--------------------------------------------------------------#
# Name: upload_backup
# Desc: upload backup files to ufile
# Arg1: backup_filepath
# Arg2: tag , backup taged with it will be removed
#--------------------------------------------------------------#
function upload_backup() {
	local backup_filepath=$1
	local tag=$2
	local filename=$(basename ${backup_filepath})

	# TODO: customize upload logic here
	log_info "[UPLOAD] upload ${backup_filepath}"
	log_info "[UPLOAD] upload to ${filename}"
	log_warn "[UPLOAD] obsolete backups: None"
	log_warn "[UPLOAD] remove obsolete backups due to retention policy"

	return 0
}


#--------------------------------------------------------------#
#                             Main                             #
#--------------------------------------------------------------#
function main() {
	# default settings
	local lock_path="/tmp/backup.lock"
	local src="postgres:///"
	local dst="/pg/backup"
	local tag=$(get_cluster_name ${src})
	local remove="1200"
	local upload="false"
	local encrypt="false"

	local filename=""
	local key=""
	local provided_filename=""
	local provided_key=""

	# parse arguments
	while (($# > 0)); do
		case "$1" in
		-s | --src=* | --url=*)
			[ "$1" == "-s" ] && shift
			src=${1##*=}
			shift
			;;
		-d | --dst=* | --dir=*)
			[ "$1" == "-d" ] && shift
			dst=${1##*=}
			shift
			;;
		-f | --file=*)
			[ "$1" == "-f" ] && shift
			provided_filename=${1##*=}
			shift
			;;
		-r | --remove=*)
			[ "$1" == "-r" ] && shift
			remove=${1##*=}
			shift
			;;
		-k | --key=*)
			[ "$1" == "-k" ] && shift
			provided_key=${1##*=}
			shift
			;;
		-t | --tag=*)
			[ "$1" == "-t" ] && shift
			tag=${1##*=}
			shift
			;;
		-u | --upload)
			upload="true"
			shift
			;;
		-e | --encrypt)
			encrypt="true"
			shift
			;;
		-h)
			usage
			exit
			;;
		*)
			usage
			exit 1
			;;
		esac
	done

	# overwrite filename & key with tag
	if [[ -z "${provided_filename}" ]]; then
		# if filename is not specified, use "backup_${tag}_${date}.tar.lz4" as filename
		filename="backup_${tag}_$(date +%Y%m%d).tar.lz4"
	else
		filename=${provided_filename}
	fi
	local backup_filepath="${dst}/${filename}"

	if [[ -z "${provided_key}" ]]; then
		# if key is not specified, use ${tag} as key
		key=${tag}
	else
		key=${provided_key}
	fi

	log_info "================================================================"
	# check parameters
	log_info "[INIT] pg-basebackup begin, checking parameters"

	if [[ ! -d ${dst} ]]; then
		log_error "[INIT] destination directory ${dst} not exist"
		exit 2
	fi

	if [[ ${remove} != [0-9]* ]]; then
		log_error "[INIT] -r,--remove should be an integer represent minutes of retention"
		exit 3
	fi

	if [[ -z $(command -v pg_basebackup) ]]; then
		log_error "[INIT] pg_basebackup binary not found in PATH"
		exit 4
	fi

	#	if [[ ${upload} == "true" ]]; then
	#		# TODO: IMPLEMENT HERE
	#	fi

	if [[ ${encrypt} == "true" ]]; then
		# if encrypt is specified, openssl sould exist
		if [[ -z $(command -v openssl) ]]; then
			log_error "[INIT] openssl binary not found in PATH when encrypt is specified"
			exit 7
		fi
	fi

	log_debug "[INIT] #====== BINARY"
	log_debug "[INIT] pg_basebackup     :   $(command -v pg_basebackup)"
	log_debug "[INIT] openssl           :   $(command -v openssl)"

	log_debug "[INIT] #====== PARAMETER"
	log_debug "[INIT] filename  (-f)    :   ${filename}"
	log_debug "[INIT] src       (-s)    :   ${src}"
	log_debug "[INIT] dst       (-d)    :   ${dst}"
	log_debug "[INIT] tag       (-t)    :   ${tag}"
	log_debug "[INIT] key       (-k)    :   ${key}"
	log_debug "[INIT] encrypt   (-e)    :   ${encrypt}"
	log_debug "[INIT] upload    (-u)    :   ${upload}"
	log_debug "[INIT] remove    (-r)    :   -mmin +${remove}"

	# Lock (Avoid multiple instance)
	if [ -e ${lock_path} ] && kill -0 $(cat ${lock_path}); then
		log_error "[LOCK] acquire lock @ ${lock_path} failed, other_pid=$(cat ${lock_path})"
		exit 8
	fi
	log_info "[LOCK] acquire lock @ ${lock_path}"
	trap "rm -f ${lock_path}; exit" INT TERM EXIT
	echo $$ >${lock_path}
	log_info "[LOCK] lock acquired success on ${lock_path}, pid=$$"

	# Start Backup
	log_info "[BKUP] backup begin, from ${src} to ${backup_filepath}"
	if [[ ${encrypt} != "true" ]]; then
		log_info "[BKUP] backup in normal mode"
		key=""
	else
		log_info "[BKUP] backup in encryption mode"
	fi

	make_backup ${src} ${backup_filepath} ${key}

	if [[ $? != 0 ]]; then
		log_error "[BKUP] backup failed!"
		exit 9
	fi
	log_info "[BKUP] backup complete!"

	# remove old local backup
	log_info "[RMBK] remove local obsolete backup: ${remove}"
	remove_backup ${dst} ${remove}
	if [[ $? != 0 ]]; then
		log_error "[RMBK] remove local obsolete backup failed!"
		exit 10
	fi
	log_info "[RMBK] remove old backup complete"

	# unlock
	rm -f ${lock_path}
	log_info "[LOCK] release lock @ ${lock_path}"

	# done
	log_info "[DONE] backup procedure complete!"
	log_info "================================================================"
}

main "$@"
